/*
 * Description: NetCard config in windows
 */

#include "NetCard.h"
#include <codecvt>
#include <comutil.h>
#include "log.h"
#include "../common/text_utils.h"

#pragma comment(lib, "wbemuuid.lib")
#pragma comment(lib, "comsuppw.lib")

#pragma execution_character_set("utf-8")

using namespace std;

void NetCard::Clear()
{
    if (pConfig) {
        pConfig.Release();
    }
    if (pObj) {
        pObj.Release();
    }
    if (pEnum) {
        pEnum.Release();
    }
    if (pService) {
        pService.Release();
    }
    if (pInst) {
        pInst.Release();
    }
    if (isInit) {
        CoUninitialize();
    }
    isInit = false;
}

bool NetCard::Init()
{
    if (isInit) {
        LOG_INFO("Init has already been completed");
        return true;
    }

    // 初始化COM接口
    HRESULT hres = CoInitializeEx(0, COINIT_MULTITHREADED);
    if (hres != S_OK) {
        return false;
    }

    // 初始化WMI调用实例
    hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID*)&pInst);
    if (hres != S_OK) {
        return false;
    }

    // 连接到本地root\cimv2命名空间,并获取 IWbemService调用的指针
    hres = pInst->ConnectServer(_bstr_t(L"ROOT\\CIMV2"), NULL, NULL, 0, NULL, 0, 0, &pService);
    if (hres != S_OK) {
        return false;
    }

    // 设置代理的安全参数
    hres = CoSetProxyBlanket(pService, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL,
        RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE);
    if (hres != S_OK) {
        return false;
    }

    if (!getNetAdapter(hres)) {
        return false;
    }
 
    VariantInit(&netPath);
    hres = pObj->Get(L"__PATH", 0, &netPath, NULL, NULL);
    if (hres != S_OK) {
        return false;
    }
    hres = pService->GetObject(_bstr_t(L"Win32_NetworkAdapterConfiguration"), 0, NULL, &pConfig, NULL);
    if (hres != S_OK) {
        return false;
    }
    isInit = true;
    LOG_INFO("Init complete");
    return true;
}

bool NetCard::getNetAdapter(HRESULT &hres)
{
    // 通过适配器名称来找到指定的适配器对象.
    CComBSTR TheQuery = L"SELECT * FROM Win32_NetworkAdapterConfiguration WHERE SettingID = \"";
    std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>, wchar_t> conversion;
    TheQuery += conversion.from_bytes(netKey).c_str();
    TheQuery += L"\"";
    hres = pService->ExecQuery(
        (BSTR)L"WQL", TheQuery, WBEM_FLAG_RETURN_IMMEDIATELY | WBEM_FLAG_FORWARD_ONLY, NULL, &pEnum);
    if (hres != S_OK) {
        return false;
    }
    ULONG num = 0;
    hres = pEnum->Next(WBEM_INFINITE, 1, &pObj, &num);
    if (hres != S_OK || num < 1) {
        TheQuery =
            L"SELECT * FROM Win32_NetworkAdapterConfiguration WHERE Description = \"Red Hat VirtIO Ethernet Adapter\"";
        hres = pService->ExecQuery(
            (BSTR)L"WQL", TheQuery, WBEM_FLAG_RETURN_IMMEDIATELY | WBEM_FLAG_FORWARD_ONLY, NULL, &pEnum);
        if (hres != S_OK) {
            return false;
        }
        num = 0;
        hres = pEnum->Next(WBEM_INFINITE, 1, &pObj, &num);
        if (hres != S_OK || num < 1) {
            LOG_ERR("Can not find netcard: {}", hres);
            return false;
        }
    }
    return true;
}

bool NetCard::EnableDhcp()
{
    LOG_INFO("Enable dhcp");
    if (!Init()) {
        return false;
    }
    return ExecuteMethod(L"EnableDHCP", NULL);
}

bool NetCard::SetIpConfig(const std::string &ip, const std::string &mask)
{
    LOG_INFO("Set ip config");
    if (!Init()) {
        return false;
    }
    
    CComPtr<IWbemClassObject> params = NULL;
    HRESULT hres = pConfig->GetMethod(_bstr_t("EnableStatic"), 0, &params, NULL);
    if (hres != S_OK) {
        return false;
    }
    CComPtr<IWbemClassObject> paramsInst = NULL;
    hres = params->SpawnInstance(0, &paramsInst);
    if (hres != S_OK) {
        return false;
    }
    auto p1 = CreateSafeArray({ ip });
    VARIANT paramVt;
    paramVt.vt = VT_ARRAY | VT_BSTR;
    paramVt.parray = p1.get();
    hres = paramsInst->Put(L"IPAddress", 0, &paramVt, NULL);
    if (hres != S_OK) {
        return false;
    }
    p1 = CreateSafeArray({ mask });
    paramVt.parray = p1.get();
    hres = paramsInst->Put(L"SubnetMask", 0, &paramVt, NULL);
    if (hres != S_OK) {
        return false;
    }
    return ExecuteMethod(L"EnableStatic", paramsInst);
}

bool NetCard::SetDns(const std::string &defaultDns, const std::string &backupDns)
{
    return SetDnsBase(false, defaultDns, backupDns);
}

bool NetCard::SetAutoDns()
{
    return SetDnsBase(true, "", "");
}

bool NetCard::SetGateway(const std::string &gateway)
{
    LOG_INFO("Set gateway");
    if (!Init()) {
        return false;
    }
    CComPtr<IWbemClassObject> params = NULL;
    HRESULT hres = pConfig->GetMethod(_bstr_t("SetGateways"), 0, &params, NULL);
    CComPtr<IWbemClassObject> paramsInst = NULL;
    hres = params->SpawnInstance(0, &paramsInst);
    if (hres != S_OK) {
        return false;
    }
    auto p1 = CreateSafeArray({ gateway });
    VARIANT paramVt;
    paramVt.vt = VT_ARRAY | VT_BSTR;
    paramVt.parray = p1.get();
    if (paramsInst->Put(L"DefaultIPGateway", 0, &paramVt, NULL) != S_OK) {
        return false;
    }
    return ExecuteMethod(L"SetGateways", paramsInst);
}

std::shared_ptr<SAFEARRAY> NetCard::CreateSafeArray(const std::vector<std::string> &args)
{
    SAFEARRAY *psa = SafeArrayCreateVector(VT_BSTR, 0, args.size());
    long idx[] = { 0 };
    for (int i = 0; i < args.size(); i++) {
        std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>, wchar_t> conversion;
        BSTR parameter = SysAllocString(conversion.from_bytes(args[i]).c_str());
        idx[0] = i;
        HRESULT hres = SafeArrayPutElement(psa, idx, parameter);
        if (hres != S_OK) {
            LOG_ERR("Failed to construct safe array {}", hres);
        }
        SysFreeString(parameter);
    }
    return shared_ptr<SAFEARRAY>(psa, [](SAFEARRAY *psa) {SafeArrayDestroy(psa); });
}

bool NetCard::SetDnsBase(bool isAuto, const std::string &defaultDns, const std::string &backupDns)
{
    LOG_INFO("Set dns");
    if (!Init()) {
        return false;
    }
    CComPtr<IWbemClassObject> params = NULL;
    HRESULT hres = pConfig->GetMethod(_bstr_t("SetDNSServerSearchOrder"), 0, &params, NULL);
    if (hres != S_OK) {
        return false;
    }
    CComPtr<IWbemClassObject> paramsInst = NULL;
    hres = params->SpawnInstance(0, &paramsInst);
    if (hres != S_OK) {
        return false;
    }
    shared_ptr<SAFEARRAY> p1;
    if (isAuto) {
        hres = paramsInst->Put(L"DNSServerSearchOrder", 0, NULL, NULL);
    } else {
        if (backupDns.size()) {
            p1 = CreateSafeArray({ defaultDns, backupDns });
        } else {
            p1 = CreateSafeArray({ defaultDns });
        }
        VARIANT paramVt;
        paramVt.vt = VT_ARRAY | VT_BSTR;
        paramVt.parray = p1.get();
        hres = paramsInst->Put(L"DNSServerSearchOrder", 0, &paramVt, NULL);
    }
    if (hres != S_OK) {
        return false;
    }
    return ExecuteMethod(L"SetDNSServerSearchOrder", paramsInst);
}

bool NetCard::ExecuteMethod(const wchar_t * method, CComPtr<IWbemClassObject> paramsInstance)
{
    bool execResult = false;
    std::string str = COMMON::TextUtils::UnicodeToUTF8(method);
    CComPtr<IWbemClassObject> results;
    HRESULT hres = pService->ExecMethod(netPath.bstrVal, _bstr_t(method), 0, NULL, paramsInstance, &results, NULL);
    if (hres != S_OK) {
        LOG_ERR("Failed to call exec method: {}", hres);
    } else {
        VARIANT vtRet;
        VariantInit(&vtRet);
        hres = results->Get(L"ReturnValue", 0, &vtRet, NULL, 0);
        if (hres != S_OK) {
            LOG_ERR("Failed to get return value: {}", hres);
        } else {
            // 0执行成功，1重启生效
            if (vtRet.uintVal == 0 || vtRet.uintVal == 1) {
                execResult = true;
            } else {
                LOG_ERR("method {} exec failed, result: {}", str, vtRet.uintVal);
            }
        }
        VariantClear(&vtRet);
    }
    return execResult;
}